import numpy as np
from major.models_manager import chat_model, embedding_model

def cos_sim(a, b):
    """计算两个向量的余弦相似度"""
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))


class SimpleRAG():

    def __init__(self):
        self.vector_db = {} # 用字典假装向量文档库
        self.chat_llm = chat_model.get_model() # 负责生成回复
        self.embedder = embedding_model.get_model() # 负责转换向量

    def add_documents(self, texts):
    # 建档：用文档建立"向量文档库"
        # 1. 用embedding模型把文档变成向量
        text_embeddings = self.embedder.embed_documents(texts)
        # 2. 把向量作为key, text作为value,保存到向量文档库
        for text, embedding in zip(texts, text_embeddings):
            self.vector_db[tuple(embedding)] = text

    # 检索：用query通过向量检索得到相关文档
    def retrieve(self, query: str, k: int = 2):
        # query 用户查询的那段字符串
        # k 返回的文档数量
        # 1. 先用embedding模型把query变成向量
        query_embedding = self.embedder.embed_query(query)
        # 2. 计算query向量和文档库中的向量之间的相似度
        sims = [(cos_sim(query_embedding, doc_vector),self.vector_db[doc_vector]) for doc_vector in self.vector_db]
        # self.vector_db[doc_vector] 返回的是文档字符串
        # 3. 根据相似度从大到小排序
        sims = sorted(sims, reverse=True,key=lambda x:x[0])
        # 4. 取前k个最相似的文档返回
        return [sim[1] for sim in sims[:k]]

    # 生成：通过得到相关文档以及query生成回复
    def generate(self, query):
        # 1. 用query去检索相关文档
        docs = self.retrieve(query)
        # 2. 用相关文档和query一起拼接成prompt
        docs_context = '\n\n'.join(docs) #文档上下文
        prompt = f'''
            请根据相关文档回答用户问题，相关文档如下, 如果文档中没有相关信息，请回答"不知道"。：
            {docs_context}
            
            用户问题：{query}
        '''
        # 3. 根据prompt生成最终回复
        res = self.chat_llm.stream(prompt)
        for chunk in res:
            print(chunk.content, end="")

# 使用示例
if __name__ == "__main__":
    # 初始化RAG系统
    rag = SimpleRAG()

    # 添加一些文档到知识库
    documents = [
        "LangChain是一个强大的开源框架，专门设计用于构建基于大语言模型(LLM)的应用程序。它提供了一系列工具和接口，使开发者能够轻松地将LLM与其他数据源和计算资源连接起来。",
        "通过LangChain，开发者可以实现复杂的自然语言处理流程，如问答系统、文本摘要和内容生成等。该框架特别擅长处理上下文管理和记忆功能，这对于构建连贯的多轮对话系统至关重要。",
        "在检索增强生成(RAG)系统中，LangChain常与向量数据库配合使用。RAG技术通过将信息检索与文本生成相结合，显著提升了生成内容的准确性和相关性。",
        "大型犬如金毛寻回犬需要每天至少90分钟的有氧运动，同时要特别注意髋关节健康；而小型犬如吉娃娃则需要更频繁但较短时间的活动，并对低温环境特别敏感。",
        "猫咪的健康管理同样复杂。现代家猫常见健康问题包括泌尿系统疾病、肥胖和牙周病。",
        "专家建议采用湿粮和干粮结合的喂养方式，每天提供新鲜流动的水源，并定期进行牙齿清洁。"
    ]
    rag.add_documents(documents) # 建档

    # 提问并获取答案
    question = "周杰伦是是谁"
    rag.generate(question)
